#coding:utf-8

# 使用labelImg标注得到的VOC格式数据集转换成COCO格式数据集
# pip install lxml
 
import os
import glob
import json
import shutil
import numpy as np
import xml.etree.ElementTree as ET
from PyQt5.QtWidgets import QMessageBox,QFileDialog
from libs.ustr import *

 
START_BOUNDING_BOX_ID = 1
 
class XmlToCoco:

    def __init__(self, labelImg, label_dict):
        self.labelImg = labelImg
        self.ImageWidth = 0
        # self.label_dict = label_dict


        # self.classes = []
        self.pre_define_categories = {}
        for label in label_dict:
            labelObj = label_dict[label] 
            self.pre_define_categories[labelObj.get_labelEn()] = labelObj.get_id()
        
        # for i, cls in enumerate(self.classes):
        #     self.pre_define_categories[cls] = i + 1
        # pre_define_categories = {'a1': 1, 'a3': 2, 'a6': 3, 'a9': 4, "a10": 5}
        self.only_care_pre_define_categories = True
        # only_care_pre_define_categories = False
    
    def get(self,root, name):
        return root.findall(name)
    
    
    def get_and_check(self,root, name, length):
        vars = root.findall(name)
        if len(vars) == 0:
            raise NotImplementedError('Can not find %s in %s.'%(name, root.tag))
        if length > 0 and len(vars) != length:
            raise NotImplementedError('The size of %s is supposed to be %d, but is %d.'%(name, length, len(vars)))
        if length == 1:
            vars = vars[0]
        return vars
    
    
    def convert(self,xml_list, json_file):
        json_dict = {"images": [], "type": "instances", "annotations": [], "categories": []}
        categories = self.pre_define_categories.copy()
        bnd_id = START_BOUNDING_BOX_ID
        all_categories = {}
        for index, line in enumerate(xml_list):
            # print("Processing %s"%(line))
            xml_f = str(line)
            img_f = xml_f.replace('annotations', 'train')
            img_f = img_f[:-4] + ".jpg"
            if not os.path.exists(img_f):
                print('文件' + xml_f + ' 没有对应的图片文件:' + img_f)
                continue

            tree = ET.parse(xml_f)
            root = tree.getroot()
            
            image_id = 20190000001 + index
            size = self.get_and_check(root, 'size', 1)
            width = int(self.get_and_check(size, 'width', 1).text)
            height = int(self.get_and_check(size, 'height', 1).text)            
            filename = os.path.basename(xml_f)[:-4] + ".jpg"
            image = {'file_name': filename, 'height': height, 'width': width, 'id':image_id}
            json_dict['images'].append(image)
            ## Cruuently we do not support segmentation
            #  segmented = get_and_check(root, 'segmented', 1).text
            #  assert segmented == '0'
            for obj in self.get(root, 'object'):
                category = self.get_and_check(obj, 'name', 1).text
                if category in all_categories:
                    all_categories[category] += 1
                else:
                    all_categories[category] = 1
                if category not in categories:
                    if self.only_care_pre_define_categories:
                        continue
                    new_id = len(categories) + 1
                    print("[warning] category '{}' not in 'pre_define_categories'({}), create new id: {} automatically".format(category, self.pre_define_categories, new_id))
                    categories[category] = new_id
                category_id = categories[category]
                bndbox = self.get_and_check(obj, 'bndbox', 1)
                xmin = int(float(self.get_and_check(bndbox, 'xmin', 1).text))
                ymin = int(float(self.get_and_check(bndbox, 'ymin', 1).text))
                xmax = int(float(self.get_and_check(bndbox, 'xmax', 1).text))
                ymax = int(float(self.get_and_check(bndbox, 'ymax', 1).text))
                assert(xmax > xmin), "xmax <= xmin, {}".format(line)
                assert(ymax > ymin), "ymax <= ymin, {}".format(line)
                o_width = abs(xmax - xmin)
                o_height = abs(ymax - ymin)
                ann = {'area': o_width*o_height, 'iscrowd': 0, 'image_id':
                    image_id, 'bbox':[xmin, ymin, o_width, o_height],
                    'category_id': category_id, 'id': bnd_id, 'ignore': 0,
                    'segmentation': []}
                json_dict['annotations'].append(ann)
                bnd_id = bnd_id + 1
    
        for cate, cid in categories.items():
            cat = {'supercategory': 'none', 'id': cid, 'name': cate}
            json_dict['categories'].append(cat)
        json_fp = open(json_file, 'w')
        json_str = json.dumps(json_dict)
        json_fp.write(json_str)
        json_fp.close()
        print("------------create {} done--------------".format(json_file))
        print("find {} categories: {} -->>> your pre_define_categories {}: {}".
            format(len(all_categories), all_categories.keys(), len(self.pre_define_categories), self.pre_define_categories.keys()))

        print("category: id --> {}".format(categories))
        print(categories.keys())
        print(categories.values())
    
    def scanXMLToCocoFun(self):
        ret = QMessageBox.question(self.labelImg, '提示', "确定读取XML文件生成Coco格式数据集吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if ret == QMessageBox.No:
            return 

        path = '.'
        dirpath = ustr(QFileDialog.getExistingDirectory(self.labelImg,
                                                       "选择要替换中文标签的XML文件路径", path,  QFileDialog.ShowDirsOnly
                                                       | QFileDialog.DontResolveSymlinks))

        if not (dirpath.endswith("annotations") or dirpath.endswith('Annotations')):
            QMessageBox.information(self.labelImg, '提示', "必须选择数据集里保存XML文件的 annotations 目录")
            return

        if dirpath is not None and len(dirpath) > 1:
            ret = QMessageBox.question(self.labelImg, '提示', "确定要扫描目录【"+dirpath+"】里的XML文件吗？",
                                       QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
            if ret == QMessageBox.Yes:
                self.toCoco(dirpath)
                return
    
    def toCoco(self, dirpath):
        dirpath = str(dirpath)
        pos = dirpath.rfind("/")
        path2 = dirpath[0: pos] + "/coco"

        train_ratio = 0.9
        
        if os.path.exists(path2 + "/annotations"):
            shutil.rmtree(path2 + "/annotations")
        os.makedirs(path2 + "/annotations")
        
        save_json_train = path2 + '/annotations/instances_train2014.json'
        save_json_val = path2 + '/annotations/instances_val2014.json'
        # xml_dir = "./tmp_xml"
    
        xml_list = glob.glob(dirpath + "/*.xml")
        for xml in xml_list:
            img = xml.replace('annotations','train')[:-4] + ".jpg"
            if not os.path.exists(img):
                os.remove(xml)
                xml_list.remove(xml)

        xml_list = np.sort(xml_list)
        np.random.seed(100)
        np.random.shuffle(xml_list)
    
        train_num = int(len(xml_list)*train_ratio)
        xml_list_train = xml_list[:train_num]
        xml_list_val = xml_list[train_num:]

        self.convert(xml_list_train, save_json_train)
        self.convert(xml_list_val, save_json_val)    

        if os.path.exists(path2 + "/images/train2014"):
            shutil.rmtree(path2 + "/images/train2014")
        os.makedirs(path2 + "/images/train2014")
        if os.path.exists(path2 + "/images/val2014"):
            shutil.rmtree(path2 +"/images/val2014")
        os.makedirs(path2 + "/images/val2014")
    
        f1 = open(path2 + "/annotations/train.txt", "w")
        print("保存train.txt到： " + path2 + '/annotations/train.txt')
        for xml in xml_list_train:
            img = xml.replace('annotations','train')[:-4] + ".jpg"
            f1.write(os.path.basename(xml)[:-4] + "\n")
            shutil.copyfile(img, path2 + "/images/train2014/" + os.path.basename(img))
    
        f2 = open(path2 + '/annotations/test.txt', "w")
        print("保存test.txt到： " + path2 + '/annotations/test.txt')
        for xml in xml_list_val:
            img = xml.replace('annotations','train')[:-4] + ".jpg"
            f2.write(os.path.basename(xml)[:-4] + "\n") 
            shutil.copyfile(img, path2 + "/images/val2014/" + os.path.basename(img))
        f1.close()
        f2.close()
        print("-------------------------------")
        print("train number:", len(xml_list_train))
        print("val number:", len(xml_list_val))